#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Jun  7 15:29:11 2023

@author: jcfq2
"""

import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import matplotlib.colors as colours

import numpy as np
# import shapely.geometry as sgeom
# from cartopy.feature import ShapelyFeature
from scipy.io import readsav
import cartopy.crs as ccrs
import sys
import os
sys.path.insert(0, '../../planetary_values')
# from grille_vogt_localtime_dawndusk_all import plot_vogt_slices
from grille_vogt_localtime_dawndusk_all import load_vogt_slices
from grille_vogt_localtime_dawndusk_all import vogt_ofl
from grille_vogt_localtime_dawndusk_all import fit_fieldline_fixed
from grille_vogt_localtime_dawndusk_all import fit_fieldline_scattered
from planet_magfield import Bmag

from shift_magnetic import to_magnetic
from shift_magnetic import from_magnetic
sys.path.insert(0, '../../tss')
from basic.polygonmask import maskeqvalue
from basic.polygonmask import masksubarray
from basic.interpolate_irregular import data_atlatlong

from astropy.modeling import models, fitting

sys.path.insert(0, '../../planetary_values')

# In[4]

colors2 = plt.cm.gist_rainbow_r(np.linspace(0.25, 1, 220))
colors1 = plt.cm.gnuplot2(np.linspace(0, 0.23, 32))

# combine them and build a new colormap
colors = np.vstack((colors1, colors2))

# colors[125:131]=[0.5,0.5,0.5,0.8]
mymap = colours.LinearSegmentedColormap.from_list('my_colormap', colors)

# In[0]
# set up values


maxint=1.0
minint=0.2
mincount=0
maxcount=3600*100.#3.*3600#hours - many hours due to bleeding of individual pixels across multiple bins
mindate=10
maxdate=30


minrange=0.7
maxrange=1.9
minrange=0.0003
maxrange=0.0018;0.0001
xxrange=6
startx=1

polelat=75
polelong=180


# In[1]
# load data


results='/Users/jcfq2/OneDrive - Northumbria University - Production Azure AD/results/'
planetary='/Users/jcfq2/OneDrive - Northumbria University - Production Azure AD/python/planetary_values/'


idl_file = readsav(results+'jupimage_phase_total_crazytown.idl')
# idl_file = readsav(results+'jupimage_phase_total_crazytown_avgclicks.idl')
# idl_file = readsav(results+'jupimage_phase_total_crazytown_avgclicks.idl')
a=idl_file['a']
a2=idl_file['a2']
b=idl_file['b']
b2=idl_file['b2']

aa=a
aa2=a2
bb=b
bb2=b2


ab=a/b
ab=np.nan_to_num(ab)


# In[4] 

totalmap=np.sum(ab,axis=0)
# plt.imshow(totalmap)
# plt.show()

totalmapN=totalmap
totalmapS=totalmap
totalmapN[0:70,:]=0
totalmapS[90:,:]=0

totalmap=totalmap/np.mean(totalmap)
totalmapN=totalmapN/np.mean(totalmapN)
totalmapS=totalmapS/np.mean(totalmapS)

lon = np.linspace(360, 0, 360)
lat = np.linspace(-90, 90, 181)
lon2d, lat2d = np.meshgrid(lon, lat)


crs = ccrs.RotatedPole(globe=ccrs.Globe(flattening=(0.0)))


# fig = plt.figure(figsize=(10, 2),dpi=300)
# for i in range(startx,xxrange):

fig = plt.figure()
# fig , (ax1 ,ax2,ax3,ax4,ax5) = plt.subplots(1, 5, gridspec_kw={'width_ratios': [1,0.2,1,1,0.2]})
# fig , (ax1 ,ax2,ax3) = plt.subplots(1, 3, gridspec_kw={'width_ratios': [1,1,1.1]})
# plt.subplots_adjust(left=0.01,
#                     bottom=0.1,
#                     right=0.9,
#                     top=0.9,
#                     wspace=0.2,
#                     hspace=0.15)

fig.set_figwidth(8.5)
fig.set_figheight(5.4)
fig.set_dpi(100)


ax1 = plt.subplot(2,2,1,projection=ccrs.Orthographic(360-180, 80))

# ax2 = plt.subplot(1,5,2)
# ax2.axis('off')

ax3 = plt.subplot(2,2,2,projection=ccrs.Orthographic(360-180, 80))
#     mapfull=ab

# ax5 = plt.subplot(1,5,5)
# ax5.axis('off')

mapfull=ab

mapmapT=np.sum(mapfull[75+1*30:75+(5+1)*30,:,:],axis=0)


gl = ax1.gridlines(crs=ccrs.PlateCarree(), linewidth=1, color='black', alpha=0.5, linestyle='dotted', draw_labels=True)
gl.xlocator = mticker.FixedLocator(np.arange(18)*20-180)
gl.ylocator = mticker.FixedLocator(np.arange(18)*10-90)
ax1.set_extent([140, 220, 50, 90], crs=ccrs.PlateCarree())


gl2 = ax3.gridlines(crs=ccrs.PlateCarree(), linewidth=1, color='black', alpha=0.5, linestyle='dotted', draw_labels=True)
gl2.xlocator = mticker.FixedLocator(np.arange(18)*20-180)
gl2.ylocator = mticker.FixedLocator(np.arange(18)*10-90)
ax3.set_extent([140, 220, 50, 90], crs=ccrs.PlateCarree())



sector1_lon,sector1_lat,sector1_loc,sector1_dis, sector2_lon,sector2_lat,sector2_loc,sector2_dis, sector3_lon,sector3_lat,sector3_loc,sector3_dis, sector4_lon,sector4_lat,sector4_loc,sector4_dis, sector5_lon,sector5_lat,sector5_loc,sector5_dis = load_vogt_slices()



ynum=3
ypos=2

ypos=(ypos-1)*6

# fig = plt.figure(figsize=(12, 2),dpi=300)
cml=180
js=80



   
ofl_lat,ofl_lon=vogt_ofl()
    
subsolarlat=0
subsolarlongitude=180#240.3#+180


s_ssl=str(int(subsolarlongitude/10+1)*10)

# print(s_ssl)

sys.path.insert(0, '../../planetary_values')

v = open(os.path.join(sys.path[0],'fieldline_tracing/fieldline_tracing_mapping_contours_kk2009ext_jrm09int_north_sslong'+s_ssl+'.txt'))
north_plotting = np.zeros([28,36,4])

header = v.readline()

for z in range(28):
    for y in range(36):
        line = v.readline()
        north_plotting[z,y,0]=line[0:12]
        north_plotting[z,y,1]=line[14:20]
        null1=line[21:29]
        if any(chr.isdigit() for chr in null1):
            north_plotting[z,y,2]=null1
        else:
            north_plotting[z,y,2]=np.nan
        null2=line[30:38]
        if any(chr.isdigit() for chr in null2):
            north_plotting[z,y,3]=null2
        else:
            north_plotting[z,y,3]=np.nan
                              
    

    

# cml=360-(75+i*30+15)    
cml=180
# 360-(90+i*20+10)   


subsolarlat=0
subsolarlongitude=cml#240.3#+180


s_ssl=str(int(subsolarlongitude/10+1)*10)

# print(s_ssl)


v = open(planetary+'fieldline_tracing/fieldline_tracing_mapping_contours_kk2009ext_jrm09int_north_sslong'+s_ssl+'.txt')
north_plotting = np.zeros([28,36,4])

header = v.readline()

for z in range(28):
    for y in range(36):
        line = v.readline()
        north_plotting[z,y,0]=line[0:12]
        north_plotting[z,y,1]=line[14:20]
        null1=line[21:29]
        if any(chr.isdigit() for chr in null1):
            north_plotting[z,y,2]=null1
        else:
            north_plotting[z,y,2]=np.nan
        null2=line[30:38]
        if any(chr.isdigit() for chr in null2):
            north_plotting[z,y,3]=null2
        else:
            north_plotting[z,y,3]=np.nan



 
np_lon2=north_plotting[1,:,3]
np_lat2=north_plotting[1,:,2]
np_loc2=north_plotting[1,:,1]

np_lon3=north_plotting[13,:,3]
np_lat3=north_plotting[13,:,2]
np_loc3=north_plotting[13,:,1]




fit_xmo,fit_ymo=fit_fieldline_fixed(np_lon2,np_lat2) 

fitx,fity=fit_fieldline_fixed(np_lon3,np_lat3) 



# fity,fitx = fit_fieldline_scattered(ofl_lon,ofl_lat,sigma=6)

plt.plot(fitx,fity,transform=ccrs.PlateCarree(),ls='--',lw='1.2')


mapmapT2,lon2d2,lat2d2=masksubarray(lon2d,lat2d,mapmapT,fit_xmo,fit_ymo)

mapmapT3,lon2d3,lat2d3=masksubarray(lon2d2,lat2d2,mapmapT2,fitx,fity,invert=True)





mapmapT4=mapmapT3[lat2d3<75]
lon2d4=lon2d3[lat2d3<75]
lat2d4=lat2d3[lat2d3<75]




cc=ax1.tricontourf(np.ravel(lon2d), np.ravel(lat2d), np.ravel(mapmapT)/np.mean(mapmapT4),transform=ccrs.PlateCarree(),levels=256,cmap='afmhot') 
ax1.tricontour(np.ravel(lon2d), np.ravel(lat2d), np.ravel(mapmapT)/np.mean(mapmapT4),transform=ccrs.PlateCarree(),levels=256,cmap='afmhot',linewidth=1) 

ax1.fill(np.append(np.arange(70)*5+5,np.array([355,5])), np.append(np.zeros(70)+80,np.array([90,90])),transform=ccrs.PlateCarree(), facecolor="none", hatch='XX', edgecolor="grey", linewidth=0.0,zorder=4)
# ax1.fill(np.append(np.arange(70)*5+5,np.array([355,5]))/360*24, np.append(np.zeros(70)+80,np.array([90,90])), facecolor="none", hatch='XX', edgecolor="grey", linewidth=0.0)
ax1.fill(north_plotting[13,:,3],north_plotting[13,:,2],transform=ccrs.PlateCarree(), facecolor="none", hatch="///", edgecolor="k", linewidth=0.0,zorder=3)
ax1.fill(np.append(np.array([350,350,270,180,90,10,10]),north_plotting[1,:,3]), np.append(np.array([90,40,40,40,40,40,90]),north_plotting[1,:,2]),transform=ccrs.PlateCarree(), facecolor="none", hatch="//", edgecolor="k", linewidth=0.0,zorder=3)

ax1.text(0.21, -0.08, '200$^{\circ}$W' ,transform=ax1.transAxes, ha="center", va="center")
ax1.text(0.5, -0.08, '180$^{\circ}$W' ,transform=ax1.transAxes, ha="center", va="center")
ax1.text(0.79, -0.08, '160$^{\circ}$W' ,transform=ax1.transAxes, ha="center", va="center")


gl.ylabels_right = False
gl2.ylabels_right = False
gl.xlabels_top = False
gl.xlabels_bottom = False
gl2.xlabels_top = False
gl2.xlabels_bottom = False


# fig.colorbar(cc,ax=ax1,location='left')


B=Bmag(lon2d4, lat2d4,radial=False)
B2=Bmag(lon2d2, lat2d2)


ax2 = plt.subplot(2,1,2)

ax2b = fig.add_axes([0.2, 0.22, 0.25, 0.025])

ax2.scatter(B,mapmapT4/np.mean(mapmapT4),marker='.',linewidth=0.1,cmap=mymap,c=360-lon2d4,vmin=90,vmax=270)
# plt.scatter(B2,mapmapT2,marker='.',linewidth=0.01)
X,Y = np.mgrid[-90:90:8,0:360:1]

# ax2.imshow(range(6)[::-1], color='green')
ax2b.imshow(Y,cmap=mymap)
# ax2b.set_xlabel('West longitude')
ax2b.set_xticks([0,180,360],['90$^{\circ}$W','180$^{\circ}$W','270$^{\circ}$W'])


ax2b.tick_params(
    axis='y',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    left=False,      # ticks along the bottom edge are off
    right=False,         # ticks along the top edge are off
    labelleft=False) # labels along the bottom edge are off


a_values=np.empty([3])
sigma_1d=np.empty([3])


g_init = models.Polynomial1D(degree=1)
#g_init = models.Gaussian1D()+ models.Polynomial1D(degree=2)
fit_g = fitting.LevMarLSQFitter()
g = fit_g(g_init, B, mapmapT4,maxiter=1000)


param0=g.parameters[0]
param1=g.parameters[1]



ax2.plot(B, g(B)/np.mean(mapmapT4),c='red',ls='--',alpha=0.5,lw='3')
ax2.set_ylabel('Norm. Intensity')
ax2.set_xlabel('|B| (Gauss)')



print(param0,param1)

B1=Bmag(lon2d, lat2d,radial=False)



scaleint=B1*param1+param0
scaleint=scaleint#/np.mean(scaleint)

dd=ax3.tricontourf(np.ravel(lon2d), np.ravel(lat2d), np.ravel(mapmapT/scaleint),transform=ccrs.PlateCarree(),levels=128,cmap='afmhot',vmin=0.1,vmax=1.4) 
ax3.tricontour(np.ravel(lon2d), np.ravel(lat2d), np.ravel(mapmapT/scaleint),transform=ccrs.PlateCarree(),levels=128,cmap='afmhot',vmin=0.1,vmax=1.4,linewidth=1) 

ax3.text(0.21, -0.08, '200$^{\circ}$W' ,transform=ax3.transAxes, ha="center", va="center")
ax3.text(0.5, -0.08, '180$^{\circ}$W' ,transform=ax3.transAxes, ha="center", va="center")
ax3.text(0.79, -0.08, '160$^{\circ}$W' ,transform=ax3.transAxes, ha="center", va="center")

cbar=fig.colorbar(cc,ax=ax3, ticks=[0.0,0.2,0.4,0.6,0.8,1.0,1.2,1.4],location="right", pad=0.08)
cbar.ax.set_ylabel('Auroral intensity')
cbar.ax.yaxis.set_label_position("right")

# fig.colorbar(dd,ax=ax5,location='left')

ax1.plot(north_plotting[1,:,3],north_plotting[1,:,2],transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='k')
ax1.plot(north_plotting[1,:,3],north_plotting[1,:,2],transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='grey',linestyle='--')
ax1.plot(north_plotting[13,:,3],north_plotting[13,:,2],transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='k')
ax1.plot(north_plotting[13,:,3],north_plotting[13,:,2],transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='grey',linestyle='--')
ax1.plot(np.arange(18)*10+90,np.zeros(18)+78,transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='k')
ax1.plot(np.arange(18)*10+90,np.zeros(18)+78,transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='grey',linestyle='--')

ax3.plot(north_plotting[1,:,3],north_plotting[1,:,2],transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='k')
ax3.plot(north_plotting[1,:,3],north_plotting[1,:,2],transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='grey',linestyle='--')
ax3.plot(north_plotting[13,:,3],north_plotting[13,:,2],transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='k')
ax3.plot(north_plotting[13,:,3],north_plotting[13,:,2],transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='grey',linestyle='--')
ax3.plot(np.arange(18)*10+90,np.zeros(18)+78,transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='k')
ax3.plot(np.arange(18)*10+90,np.zeros(18)+78,transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='grey',linestyle='--')




print(north_plotting[13,:,0])
print(param0,param1)

fig.savefig('asd_mag.pdf', dpi=300, bbox_inches='tight', facecolor='white') 


# %%


# blevels = np.arange(12)*2

# fig = plt.figure()

# # B2 = B1+0

# # step_size = 1
# # maxB = np.max(B2)

# # for i in range(int((maxB+0.5)/step_size)):
# #     B2[B1>step_size*i]=step_size*i
# #     print(step_size*i)

# fig.set_figwidth(8.5/2)
# fig.set_figheight(5.4/2)
# fig.set_dpi(100)


# ax1 = plt.subplot(1,1,1,projection=ccrs.Orthographic(360-180, 80))
# # cc = ax1.tricontourf(np.ravel(lon2d), np.ravel(lat2d), np.ravel(B1),transform=ccrs.PlateCarree(),cmap='rainbow',linewidth=1,vmin=0) 

# # ax1.tricontour(np.ravel(lon2d), np.ravel(lat2d), np.ravel(B1),transform=ccrs.PlateCarree(),cmap='rainbow',linewidth=1) 
# cc=ax1.imshow(np.flip(np.flip(B1,axis=0),axis=1),transform=ccrs.PlateCarree(),cmap='rainbow',extent=(360,0,-90,90) )

# ax1.plot(north_plotting[1,:,3],north_plotting[1,:,2],transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='k')
# ax1.plot(north_plotting[1,:,3],north_plotting[1,:,2],transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='grey',linestyle='--')
# ax1.plot(north_plotting[13,:,3],north_plotting[13,:,2],transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='k')
# ax1.plot(north_plotting[13,:,3],north_plotting[13,:,2],transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='grey',linestyle='--')
# ax1.plot(np.arange(18)*10+90,np.zeros(18)+78,transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='k')
# ax1.plot(np.arange(18)*10+90,np.zeros(18)+78,transform=ccrs.PlateCarree(),linewidth=1,alpha=0.5,c='grey',linestyle='--')

# gl = ax1.gridlines(crs=ccrs.PlateCarree(), linewidth=1, color='black', alpha=0.5, linestyle='dotted', draw_labels=True)
# gl.xlocator = mticker.FixedLocator(np.arange(18)*20-180)
# gl.ylocator = mticker.FixedLocator(np.arange(18)*10-90)
# ax1.set_extent([140, 220, 50, 90], crs=ccrs.PlateCarree())
# # ax1.set_extent([90, 270, 0, 90], crs=ccrs.PlateCarree())

# cbar=fig.colorbar(cc,ax=ax1,location="right", pad=0.08,ticks=[0,2,4,6,8,10,12,14,16,18,20])
# cbar.ax.set_ylabel('|B| mag')
# cbar.ax.yaxis.set_label_position("right")
# gl.ylabels_right = False
# gl2.ylabels_right = False
# gl.xlabels_top = False
# gl.xlabels_bottom = False
# gl2.xlabels_top = False
# gl2.xlabels_bottom = False


# fig.savefig('asd_onlymag.pdf', dpi=300, bbox_inches='tight', facecolor='white') 


# cc=ax1.tricontourf(np.ravel(lon2d), np.ravel(lat2d), np.ravel(B1),transform=ccrs.PlateCarree()) 
# cc = ax1.imshow(B1)
# In[7]
# ax3.plot(B, g(B)/np.mean(mapmapT4),c='red',ls='--',alpha=0.5)


print(np.corrcoef(B,mapmapT4)) # 0.9357598
# print(np.corrcoef(total_eq_bright,total_po_bright)) # 0.9357598



